from collections import OrderedDict

import torch.nn as nn

__all__ = ['xception65']

from training.models.helpers import load_pretrained

default_cfgs = {
    'xception65': {
        'url': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/gluon_xception-7015a15c.pth',
        'input_size': (3, 304, 304),
        'crop_pct': 0.903,
        'pool_size': (10, 10),
        'interpolation': 'bicubic',
        'mean': (0.485, 0.456, 0.406),
        'std': (0.229, 0.224, 0.225),
        'num_classes': 1000,
        'first_conv': 'conv1',
        'classifier': 'fc'
        # The resize parameter of the validation transform should be 333, and make sure to center crop at 299x299
    }
}


# Calculate symmetric padding for a convolution
def get_padding(kernel_size: int, stride: int = 1, dilation: int = 1, **_) -> int:
    padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
    return padding


class SeparableConv2d(nn.Module):

    def __init__(self, inplanes, planes, kernel_size=3, stride=1, dilation=1, bias=False, norm_layer=None,
                 norm_kwargs=None):
        super(SeparableConv2d, self).__init__()
        norm_kwargs = norm_kwargs if norm_kwargs is not None else {}
        self.kernel_size = kernel_size
        self.dilation = dilation

        # depthwise convolution
        padding = get_padding(kernel_size, stride, dilation)
        self.conv_dw = nn.Conv2d(inplanes, inplanes, kernel_size, stride=stride, padding=padding, dilation=dilation,
                                 groups=inplanes, bias=bias)
        self.bn = norm_layer(num_features=inplanes, **norm_kwargs)
        # pointwise convolution
        self.conv_pw = nn.Conv2d(inplanes, planes, kernel_size=1, bias=bias)

    def forward(self, x):
        x = self.conv_dw(x)
        x = self.bn(x)
        x = self.conv_pw(x)
        return x


class Block(nn.Module):
    r"""
    Args:
        planes (int or tuple): output channel
    """

    def __init__(self, inplanes, planes: (int or tuple), stride=1, dilation=1, start_with_relu=True,
                 norm_layer=None, norm_kwargs=None, ):
        super(Block, self).__init__()
        norm_kwargs = norm_kwargs if norm_kwargs is not None else {}
        if isinstance(planes, (list, tuple)):
            assert len(planes) == 3
        else:
            planes = (planes,) * 3
        outplanes = planes[-1]

        if outplanes != inplanes or stride != 1:
            self.skip = nn.Sequential()
            self.skip.add_module('conv1', nn.Conv2d(inplanes, outplanes, 1, stride=stride, bias=False)),
            self.skip.add_module('bn1', norm_layer(num_features=outplanes, **norm_kwargs))
        else:
            self.skip = None

        rep = OrderedDict()
        for i in range(3):
            rep['act%d' % (i + 1)] = nn.ReLU(inplace=True)
            rep['conv%d' % (i + 1)] = SeparableConv2d(inplanes, planes[i], 3, stride=stride if i == 2 else 1,
                                                      dilation=dilation, norm_layer=norm_layer, norm_kwargs=norm_kwargs)
            rep['bn%d' % (i + 1)] = norm_layer(planes[i], **norm_kwargs)
            inplanes = planes[i]

        if not start_with_relu:
            del rep['act1']
        else:
            rep['act1'] = nn.ReLU(inplace=False)
        self.rep = nn.Sequential(rep)

    def forward(self, x):
        skip = x
        if self.skip is not None:
            skip = self.skip(skip)
        x = self.rep(x) + skip
        return x


class Xception65(nn.Module):
    """
    Modified Alighed Xception
    """

    def __init__(self, num_classes=1000, in_chans=3, output_stride=32, norm_layer=nn.BatchNorm2d,
                 norm_kwargs=None, drop_rate=0., global_pool='avg'):
        super(Xception65, self).__init__()
        self.num_classes = num_classes
        self.drop_rate = drop_rate
        norm_kwargs = norm_kwargs if norm_kwargs is not None else {}
        if output_stride == 32:
            entry_block3_stride = 2
            exit_block20_stride = 2
            middle_block_dilation = 1
            exit_block_dilations = (1, 1)
        elif output_stride == 16:
            entry_block3_stride = 2
            exit_block20_stride = 1
            middle_block_dilation = 1
            exit_block_dilations = (1, 2)
        elif output_stride == 8:
            entry_block3_stride = 1
            exit_block20_stride = 1
            middle_block_dilation = 2
            exit_block_dilations = (2, 4)
        else:
            raise NotImplementedError

        # Entry flow
        self.conv1 = nn.Conv2d(in_chans, 32, 3, 2, 1, bias=False)
        self.bn1 = norm_layer(num_features=32, **norm_kwargs)
        self.act1 = nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(32, 64, 3, 1, 1, bias=False)
        self.bn2 = norm_layer(num_features=64)
        self.act2 = nn.ReLU(inplace=True)

        self.block1 = Block(64, 128, 2, start_with_relu=False, norm_layer=norm_layer, norm_kwargs=norm_kwargs)
        self.block1_act = nn.ReLU(inplace=True)
        self.block2 = Block(128, 256, 2, start_with_relu=False, norm_layer=norm_layer, norm_kwargs=norm_kwargs)
        self.block3 = Block(256, 728, entry_block3_stride, norm_layer=norm_layer, norm_kwargs=norm_kwargs)

        # Middle flow
        self.mid = nn.Sequential(
            OrderedDict([
                (
                    'block%d' % i,
                    Block(728, 728, 1, dilation=middle_block_dilation, norm_layer=norm_layer, norm_kwargs=norm_kwargs)
                ) for i in range(4, 20)
            ]))

        # Exit flow
        self.block20 = Block(728, (728, 1024, 1024), stride=exit_block20_stride, dilation=exit_block_dilations[0],
                             norm_layer=norm_layer, norm_kwargs=norm_kwargs)
        self.block20_act = nn.ReLU(inplace=True)

        self.conv3 = SeparableConv2d(1024, 1536, 3, stride=1, dilation=exit_block_dilations[1],
                                     norm_layer=norm_layer, norm_kwargs=norm_kwargs)
        self.bn3 = norm_layer(num_features=1536, **norm_kwargs)
        self.act3 = nn.ReLU(inplace=True)

        self.conv4 = SeparableConv2d(1536, 1536, 3, stride=1, dilation=exit_block_dilations[1],
                                     norm_layer=norm_layer, norm_kwargs=norm_kwargs)
        self.bn4 = norm_layer(num_features=1536, **norm_kwargs)
        self.act4 = nn.ReLU(inplace=True)

        self.num_features = 2048
        self.conv5 = SeparableConv2d(1536, self.num_features, 3, stride=1, dilation=exit_block_dilations[1],
                                     norm_layer=norm_layer, norm_kwargs=norm_kwargs)
        self.bn5 = norm_layer(num_features=self.num_features, **norm_kwargs)
        self.act5 = nn.ReLU(inplace=True)
        self.feature_info = [
            dict(num_chs=64, reduction=2, module='act2'),
            dict(num_chs=128, reduction=4, module='block1_act'),
            dict(num_chs=256, reduction=8, module='block3.rep.act1'),
            dict(num_chs=728, reduction=16, module='block20.rep.act1'),
            dict(num_chs=2048, reduction=32, module='act5'),
        ]

    def forward_features(self, x):
        # Entry flow
        x = self.conv1(x)
        x = self.bn1(x)
        x = self.act1(x)

        x = self.conv2(x)
        x = self.bn2(x)
        x = self.act2(x)

        x = self.block1(x)
        x = self.block1_act(x)
        c1 = x
        x = self.block2(x)
        # c2 = x
        x = self.block3(x)

        # Middle flow
        x = self.mid(x)
        # c3 = x

        # Exit flow
        x = self.block20(x)
        x = self.block20_act(x)
        x = self.conv3(x)
        x = self.bn3(x)
        x = self.act3(x)

        x = self.conv4(x)
        x = self.bn4(x)
        x = self.act4(x)

        x = self.conv5(x)
        x = self.bn5(x)
        x = self.act5(x)
        return c1, x

    def forward(self, x):
        low_level_feat, high_level_feat = self.forward_features(x)
        return low_level_feat, high_level_feat


def xception65(pretrained=False, output_stride=16):
    """
    Modified Aligned Xception-65
    """
    default_cfg = default_cfgs['xception65']
    model = Xception65(output_stride=output_stride)
    model.default_cfg = default_cfg
    if pretrained:
        load_pretrained(model, default_cfg, strict=False)
    return model
